#define M_PI 3.14159

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <iostream>
#include <fstream>

//#include <Windows.h>
#include <thread>

#include "mat_utils.h"
#include "GVFField.h"
#include "snake.h"

using namespace std;
using namespace cv;

Snake::Snake(double radius, ofstream& in_logfile, GVFField& ingvf,
		double in_alpha, double in_beta, double in_kappa, double in_gamma,
		double in_mindist, double in_maxdist)  {
	currentstep = -1;
	gvf = &ingvf;
	snakefile = &in_logfile;

	ddepth = CV_64F;
	///BEGIN SNAKE CODE
	//snake constants
	//NOTE - these can be rather finicky, especially gamma
		//increasing or decreasing gamma can lead to convergence or non-convergence of snake
		//tinker with these until satisfied with the results
	alpha = in_alpha;		//snake tension, anti-stretching constant
	beta = in_beta;		//snake rigidity, anti-bending constant
	kappa = in_kappa;		//ffield coupling constant
	gamma = in_gamma;		//step size

	maxdist = in_mindist;		//furthest dist 2 snake points can be
	mindist = in_maxdist;		//closest dist 2 snake points can be
	avgdist = mindist + ((maxdist - mindist)/2.0);
	
	int numpts = int(2 * M_PI * radius / avgdist);			//number of points in snake

	pts = Mat(numpts, 2, ddepth, Scalar(0.0));
	
	

	//TEMPORARY - initialize snake in a circle around the image
	int midx = gvf->u.cols / 2;
	int midy = gvf->u.rows / 2;
	
	double dtheta = 2 * M_PI / numpts;
	for (int i=0; i<numpts; i++) {
		double theta = dtheta * i;
		pts.at<double>(i, 0) = midx + radius*cos(theta);
		pts.at<double>(i, 1) = midy + radius*sin(theta);
	}
	currentstep = 0;

	//ofstream ffieldfile;
	//ffieldfile.open("ffield.txt", ios::out);

	//snakefile.setf(std::ios::fixed, std:: ios::floatfield);
	//snakefile << "Initial Snake"<<endl<<mirror<<endl;
	
}

void Snake::deformSnake(int iter, int num_images) {
	int ddepth = CV_64F;
	int modfactor = iter/num_images;
	Mat aby, aby_inv;
	//compute (A + yI)^-1 => pentadiagonal banded matrix
	Mat ones = Mat(pts.rows, pts.rows, ddepth, Scalar(1.0));
	aby = Mat(pts.rows, pts.rows, ddepth, Scalar(0.0));
	
	//NOTE snake paper says 4*beta, but MatLab code says 6*beta, and 6*beta works...
	double mdiag = 2*alpha + 6*beta + gamma;
	double diag1 = -alpha - 4*beta;
	
	//set edge points - allows closed loop functionality
	aby.at<double>(0,aby.cols-1) = diag1;
	aby.at<double>(0,aby.cols-2) = beta;
	aby.at<double>(1,aby.cols-1) = beta;

	aby.at<double>(aby.rows-1,0) = diag1;
	aby.at<double>(aby.rows-1,1) = beta;
	aby.at<double>(aby.rows-2,0) = beta;

	for (int i=0;i<pts.rows; i++) {
		aby.at<double>(i,i) = mdiag;
		if (i < pts.rows-2) {
			aby.at<double>(i,i+2) = (beta);
			aby.at<double>(i+2,i) = (beta);
			aby.at<double>(i,i+1) = diag1;
			aby.at<double>(i+1,i) = diag1;
		}
		else if (i < pts.rows-1) {
			aby.at<double>(i,i+1) = diag1;
			aby.at<double>(i+1,i) = diag1;
		}
	}
	aby_inv = aby.inv(DECOMP_LU);
	zero_mat2d(aby_inv, 0.00001);
	//snakefile << "aby_inv"<<endl<<aby_inv<<endl;

	//relax the snake iteratively
	int maxx = gvf->u.cols-1;	//width of image
	int maxy = gvf->u.rows-1;	//height of image
	for (int i=0; i<iter; i++) {
		Mat snake_ffield(pts.rows, 2, ddepth);		//collection of points corresponding to the force field on each point in the snake
		Mat snakeimg = Mat(maxx+1, maxy+1, ddepth, Scalar(0.0));	//temporary snake image printing
		//Mat snakeimg2 = Mat(maxx+1, maxy+1, ddepth, Scalar(0.0));	//temporary snake image printing
		//compute the new snake point locations
		for(int j=0; j<pts.rows; j++) {
//TODO compress these operations for efficiency
	//might try an interpolation method from the OpenCV lib
			//coordinates of snake point
			double x = pts.at<double>(j, 0);
			double y = pts.at<double>(j, 1);
			int ix = int(floor(x + 0.5));
			int iy = int(floor(y + 0.5));

			//keep indices within the image (bounds checking)
			if (ix < 0) {
				cout << "x[" << j << "]: " << ix << endl;
				ix = 0;
			}
			else if (ix > maxx) {
				cout << "x[" << j << "]: " << ix << endl;
				ix = maxx;
			}
			if (iy < 0) {
				cout << "y[" << j << "]: " << iy << endl;
				iy = 0;
			}
			else if (iy > maxy) {
				cout << "y[" << j << "]: " << iy << endl;
				iy = maxy;
			}

			snakeimg.at<double>(iy, ix) = 20.0;
			//compute y*x + f_x, where f_x = u,v
			double du = kappa*gvf->u.at<double>(ix,iy);
			double dv = kappa*gvf->v.at<double>(ix,iy);
			double nx = 0;
			double ny = 0;

			nx = gamma * x + du;
			ny = gamma * y + dv;

			snake_ffield.at<double>(j,0) = nx;
			snake_ffield.at<double>(j,1) = ny;
		}
		cout << "step " << i <<endl;
		pts = aby_inv * snake_ffield;

		int numpts = pts.rows;

//BEGIN SNAKE INTERPOLATION
		//redistribute the points along the snake, "snake interpolation"
		//NOTE indices 0 and 1 of the mirror are faulty due to absent neighbors

		//Recompute number of points on the snake
		int numptsNew = 0;//numptsNew + 4;
		int currentpt = 3;			//current index of the original snake
		int currentSnakept = 3;		//index of the redistrib snake being generated ~ kinda the number of points generated by the loop
		double lastdist = avgdist;		//distance from last point left to cover/walk
		
		Mat snakeNew = Mat(3, 2, ddepth);
		//sets the initial previous point for the redistributed snake
		snakeNew.at<double>(2, 0) = pts.at<double>(2, 0);
		snakeNew.at<double>(2, 1) = pts.at<double>(2, 1);

		//Redistribute the points along the snake line (ProNOI-style method)
			//this should end with currentpt=mirrorpts-2?
		int laststep = -1;		//temporary, used to check what was the last iteration action taken

		//Fix index 2 and use that to start with, previous points are bad
		while (currentpt < pts.rows-1) {
//TODO this could be simplified to 1 line, or an iterator method
			//current old snake point
			double x = pts.at<double>(currentpt, 0);
			double y = pts.at<double>(currentpt, 1);

			//Get the previous point for dist calculation and direction projection
			double px = 0;
			double py = 0;
			//if we generated a point last iteration, use it for the dist
			if (lastdist == avgdist) {
				px = snakeNew.at<double>(currentSnakept-1, 0);
				py = snakeNew.at<double>(currentSnakept-1, 1);
			}
			//else use the next old snake point
			else {
				px = pts.at<double>(currentpt-1, 0);
				py = pts.at<double>(currentpt-1, 1);
			}
			
			//calculate distance between the points
			double dx = x - px;
			double dy = y - py;
			double dist = sqrt(pow(dx,2) + pow(dy,2));

			//add point on line going towards next point (line projection) and reset lastdist
			if (dist > lastdist) {
				double nx = px + (dx / dist * lastdist);
				double ny = py + (dy / dist * lastdist);
				Mat newpt = Mat (1, 2, ddepth);
				newpt.at<double>(0, 0) = nx;
				newpt.at<double>(0, 1) = ny;
				snakeNew.push_back(newpt);
				currentSnakept++;
				lastdist = avgdist;
				laststep = 0;

				//if (nx >0 && ny > 0 && nx < snakeimg.rows && ny < snakeimg.cols) snakeimg2.at<double>(nx, ny) = 20.0;
			}
			//otherwise, move to next point
			else if (dist < lastdist)
			{
				currentpt++;
				lastdist = lastdist - dist;
				laststep = 1;
			}
			//they're equal, weird
			//add the point at the current snake position and reset lastdist
			else {
				snakeNew.at<double>(currentSnakept, 0) = x;
				snakeNew.at<double>(currentSnakept, 1) = y;
				lastdist = avgdist;
				currentpt++;
				currentSnakept++;
				laststep = 2;
				//if (x >0 && y > 0 && x < snakeimg.rows && y < snakeimg.cols) snakeimg2.at<double>(x, y) = 20.0;
			}
		}
		
		//here, it might be currentpt-1 was the last point used
			//or currentSnakept-1 was the last point set,
			//or both?
		//set head and tail points for the new snake
		//add last 2 points as the mirror of the first 2
		Mat newpts = Mat(2, 2, ddepth);
		newpts.at<double>(0,0) = snakeNew.at<double>(2,0);
		newpts.at<double>(0,1) = snakeNew.at<double>(2,1);
		newpts.at<double>(1,0) = snakeNew.at<double>(3,0);
		newpts.at<double>(1,1) = snakeNew.at<double>(3,1);
		snakeNew.push_back(newpts);

		snakeNew.at<double>(0,0) = snakeNew.at<double>(snakeNew.rows-4,0);
		snakeNew.at<double>(0,1) = snakeNew.at<double>(snakeNew.rows-4,1);
		snakeNew.at<double>(1,0) = snakeNew.at<double>(snakeNew.rows-3,0);
		snakeNew.at<double>(1,1) = snakeNew.at<double>(snakeNew.rows-3,1);


		pts = snakeNew;
		//NOTE all point vals are now INVALID/OLD
		numptsNew = pts.rows;
		//recompute aby_inv if the number of points changes
//TODO optimize this by simply adding rows and columns
//OR compute a handful of aby_inv values for different sizes and store in memory
		if (numpts != numptsNew) {
			numpts = numptsNew;
			aby = Mat(pts.rows, pts.rows, ddepth, Scalar(0.0));
			//set edge points - allows closed loop functionality
			aby.at<double>(0,aby.cols-1) = diag1;
			aby.at<double>(0,aby.cols-2) = beta;
			aby.at<double>(1,aby.cols-1) = beta;

			aby.at<double>(aby.rows-1,0) = diag1;
			aby.at<double>(aby.rows-1,1) = beta;
			aby.at<double>(aby.rows-2,0) = beta;

			for (int j=0;j<pts.rows; j++) {
				//NOTE snake paper says 4*beta, but MatLab code says 6*beta, and 6*beta works...
				aby.at<double>(j,j) = mdiag;
				if (j < pts.rows-2) {
					aby.at<double>(j,j+2) = beta;
					aby.at<double>(j+2,j) = beta;
					aby.at<double>(j,j+1) = diag1;
					aby.at<double>(j+1,j) = diag1;
				}
				else if (j < pts.rows-1) {
					aby.at<double>(j,j+1) = diag1;
					aby.at<double>(j+1,j) = diag1;
				}
			}
			aby_inv = aby.inv(DECOMP_LU);
			zero_mat2d(aby_inv, 0.00001);
			//print_mat2d(aby_inv, "aby.txt", ", ");
		}
//END SNAKE INTERPOLATION


		currentstep++;
		if (i%modfactor == 0) {
			//snakefile << "SNAKE[" << i << "]:" <<endl<<pts<<endl;
			char numstr[5];
			_itoa(i, numstr, 10);
			string winstr = "snake";
			winstr = winstr.append(numstr);
			//namedWindow(winstr, WINDOW_NORMAL);
			//imshow(winstr, snakeimg);
			imgData *snakeimgData = new imgData(numstr, snakeimg.clone());
			std::thread dispThread(dispImg, snakeimgData);
			dispThread.detach();
			//TEMP show adjusted snake after reordering points
			//string winstr2 = "snAdjusted";
			//winstr2.append(numstr);
			//namedWindow(winstr2, WINDOW_NORMAL);
			//imshow(winstr2, snakeimg2);

			*(snakefile) << *this;
		}
		//ffieldfile.close();
	}
	*(snakefile) << *this;
}


std::ostream &operator<<(std::ostream &pipe, Snake &out) {
	pipe <<"Step " << out.currentstep << endl;

	int numpts = out.pts.rows;
	double *p;
	int i = 0;
	for (i = 0; i < numpts; i++) {
		p = out.pts.ptr<double>(i);
		pipe << *p << " " << *(p + 1) << endl;
	}
	return pipe;
}
std::istream &operator>>(std::istream &pipe, Snake &in) {
	//TODO implement this later when more info on reboot process is available
	return pipe;
}